from typing import Union,Optional,Iterator,TypedDict,Any
from pathlib import Path

from langchain_community.document_loaders import DirectoryLoader,TextLoader
from langchain_core.prompt_values import PromptValue
from langchain_core.documents import Document
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import BaseMessage
from langchain_text_splitters import RecursiveCharacterTextSplitter 
from langchain.chat_models import init_chat_model
from langchain import hub
from langchain.prompts import PromptTemplate,ChatPromptTemplate
from langchain_mongodb import MongoDBAtlasVectorSearch
from langchain_openai import OpenAIEmbeddings

from langgraph.graph import START,StateGraph
from langgraph.graph.state import CompiledStateGraph

from model_config import get_embedding_model_default,get_chat_model_default
from vector_store_config import get_vector_storage_default

# 加载模型
llm:BaseChatModel=get_chat_model_default()
# 嵌入模型
embedding=get_embedding_model_default()
# 向量存储
vector_store=get_vector_storage_default(embedding=embedding)

prompt:PromptTemplate=ChatPromptTemplate([
  ('human','''
   下面的上下文部分是订单信息，格式参考这个粒子：
<订单信息开始>
订单id:1137,
订单编号:HY11924286,
订单状态:4,
下单时间:2023-07-05 11:39:31,
订单金额:128.8,
订单地址:河北省沧州市沧县河北采油三厂畅心公寓(订单号:HY11924286),
订单收件人:马庆国,
订单联系电话:13389017758,
订单留言:天津分公司-南部项目部办公用品,
下单企业:中国石油集团测井有限公司天津分公司,
企业联系人:123,
企业联系电话:123,
企业地址:123,
<订单商品信息开始>

商品1:
  商品名称:"+得力 长尾夹 8555ES 5#19mm 彩色长尾夹票夹40只/盒 一盒起订"
  购买数量:10,
  商品单价:11.04,
  商品合计金额:110.4,
  商品规格:8555ES 5#19mm,
  商品单位:个,
  商品品牌:得力,
  商品产地:中国浙江,


商品2:
  商品名称:"+得力(deli) 长尾夹 3#中号 32mm 彩色长尾夹票夹24个/盒"
  购买数量:1,
  商品单价:18.4,
  商品合计金额:18.4,
  商品规格:3#中号 32mm,
  商品单位:盒,
  商品品牌:得力(deli),
  商品产地:中国,

<订单商品信息结束>
<订单信息结束>   
   你是问答任务的助手。使用以下检索到的上下文来回答问题。如果你不知道答案，就说你不知道。最多使用三句话并保持答案简洁。
   问题:{question}
   上下文:{context}
   '''),
])

class State(TypedDict):
  question:str
  context:list[Document]
  answer:str

def retrieve(state:State)->dict[str,list[Document]]:
  retrieved_docs:list[Document]=vector_store.similarity_search(state['question'],k=10)
  return {'context':retrieved_docs}

def generate(state:State)->dict[str,str|list[str|dict]]:
  docs_content:str='\n\n'.join(doc.page_content for doc in state['context'])
  print(docs_content)
  messages:PromptValue=prompt.invoke({'question':state['question'],'context':docs_content})
  response:BaseMessage=llm.invoke(messages)
  return {'answer':response.content}

graph_builder:StateGraph=StateGraph(State).add_sequence([retrieve,generate])
graph_builder.add_edge(START,'retrieve')
graph:CompiledStateGraph=graph_builder.compile()

# query='谁和谁互发新年贺词'
query='新疆吐哈石油勘探开发有限公司最近购买了什么商品'
response:Union[dict[str,Any]|Any]=graph.invoke({'question':query})
print(response['answer'])
